package com.unbox.admin.util;

import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.FrameGrabber;
import org.bytedeco.javacv.Java2DFrameConverter;
import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;

/**
 * 视频获取图片
 */
public class JavaCvUtil {

    /**
     * 获取视频帧数图片
     * @param fileDatas
     * @return
     */
    public static byte[] videoToPic(byte[] fileDatas) {
        try(InputStream inputStream = new ByteArrayInputStream(fileDatas)){
            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
            BufferedImage bufferedImage = getBufferedImageByFrame(10,inputStream);
            ImageIO.write(bufferedImage, "jpg", outputStream);
            byte[] imageBytes = outputStream.toByteArray();
            outputStream.close();
            return imageBytes;
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 获取指定帧数的封面图片
     *
     * @param frameNum 帧数
     * @param inputStream 文件流
     */
    public static BufferedImage getBufferedImageByFrame(int frameNum, InputStream inputStream) throws IOException {
        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(inputStream);
        return getBufferedImageByFrame(frameNum, grabber);
    }

    private static BufferedImage getBufferedImageByFrame(int frameNum, FFmpegFrameGrabber grabber)
            throws FrameGrabber.Exception {
        grabber.start();
        Frame frame;
        int i = 0;
        int fps = (int)grabber.getFrameRate();
        BufferedImage buffer = null;
        while (i < grabber.getLengthInFrames()) {
            frame = grabber.grabImage();
            if (i >= fps && i % (fps * frameNum) == 0) {
                Java2DFrameConverter converter = new Java2DFrameConverter();
                buffer = converter.getBufferedImage(frame);
                break;
            }
            i++;
        }
        grabber.stop();
        return buffer;
    }
}
